import numpy as np
import matplotlib.pyplot as plt
from scipy.fft import fft2, ifft2, fftshift, ifftshift
from PIL import Image

# 读取彩色PNG图片
image = plt.imread('bg1.png')
h=image.shape[0]
w=image.shape[1]

# 读取签名图片，黑色内容为有效内容，有效内容越多对图片影响越大
sign=Image.open('sign.png')
sign=sign.resize((w//2,h//2))
sign=sign.im

# 分离颜色通道
red_channel = image[:, :, 0]
green_channel = image[:, :, 1]
blue_channel = image[:, :, 2]

# 对颜色通道进行傅里叶变换
f_red = fft2(red_channel)
f_green = fft2(green_channel)
f_blue = fft2(blue_channel)
# 乘的话抗干扰强，但是对图片影响大。除的话抗干扰能力小，但是基本看不出差别。
for (i,v) in enumerate(sign):
    if v[0]<50:
        f_red[i//w][i%w]*=5
        # f_green[i//w][i%w]*=3
        # f_blue[i//w][i%w]*=3

# 将零频率分量移到频谱中心
fshift_red = fftshift(f_red)
fshift_green = fftshift(f_green)
fshift_blue = fftshift(f_blue)

# 如果你想可视化频谱图，可以计算幅度谱
# magnitude_spectrum_red = 20 * np.log(np.abs(fshift_red))
# magnitude_spectrum_green = 20 * np.log(np.abs(fshift_green))
# magnitude_spectrum_blue = 20 * np.log(np.abs(fshift_blue))

# plt.imsave('magnitude_spectrum_red.png', magnitude_spectrum_red, cmap='gray')
# plt.imsave('magnitude_spectrum_green.png', magnitude_spectrum_green, cmap='gray')
# plt.imsave('magnitude_spectrum_blue.png', magnitude_spectrum_blue, cmap='gray')

# 假设你想从频谱图还原原始图像，对每个颜色通道进行逆傅里叶变换
f_ishift_red = ifftshift(fshift_red)
f_ishift_green = ifftshift(fshift_green)
f_ishift_blue = ifftshift(fshift_blue)

# 逆傅里叶变换
image_back_red = np.abs(ifft2(f_ishift_red))
image_back_green = np.abs(ifft2(f_ishift_green))
image_back_blue = np.abs(ifft2(f_ishift_blue))

# 合并颜色通道
image_back = np.dstack((image_back_red, image_back_green, image_back_blue))

# 将像素值归一化到0-255范围并转换为uint8类型
image_back = (image_back - np.min(image_back)) / (np.max(image_back) - np.min(image_back)) * 255
image_back = image_back.astype(np.uint8)

plt.imsave('image_back.png', image_back)